package com.test.springai.service;

import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Component;

/**
 * Created with IntelliJ IDEA.
 *
 * @description: 默认的查询重写器实现 负责利用 LLM 将用户原始查询重写为更适合向量数据库检索的单一查询。
 * @author: liuziyang
 * @since: 9/3/25 11:43 AM
 * @modifiedBy:
 * @version: 1.0
 */
@Component
public class DefaultQueryRewriter implements QueryRewriter {
  private final ChatClient chatClient;
  // 用于 LLM 调用的线程池
  private final Executor threadPoolTaskExecutor;
  // 提示模板：用于将原始查询重写为更适合检索的单一查询
  private final String REWRITE_SINGLE_QUERY_PROMPT =
      """
            你是一个查询优化助手。请根据以下用户问题，将其改写成一个更简洁、精确，且更适合从向量数据库中检索文档的单一查询。
            确保改写后的查询能够捕获用户问题的核心意图，并包含重要的关键词。
            不要添加任何额外解释，只返回改写后的查询字符串。
            用户原始问题: {original_query}
            """;

  /**
   * 构造函数，用于初始化 DefaultQueryRewriter。
   *
   * @param chatClientBuilder 注入的 ChatClient.Builder 用于构建 ChatClient 实例
   * @param threadPoolTaskExecutor 注入的线程池，用于异步执行 LLM 调用
   */
  public DefaultQueryRewriter(
      ChatClient.Builder chatClientBuilder,
      @Qualifier("applicationTaskExecutor") Executor threadPoolTaskExecutor) {
    this.chatClient = chatClientBuilder.build();
    this.threadPoolTaskExecutor = threadPoolTaskExecutor;
  }

  /**
   * 将用户原始查询异步重写为单个优化后的查询。 该方法会向 LLM 发送一个提示，指导其进行查询改写。
   *
   * @param originalQuery 用户的原始查询字符串。
   * @return 包含优化后查询字符串的 CompletableFuture。
   */
  @Override
  public CompletableFuture<String> rewriteSingleQuery(String originalQuery) {
    // 在指定的线程池中执行异步任务
    return CompletableFuture.supplyAsync(
        () -> {
          System.out.println("Rewriting single query for: " + originalQuery);
          // 1. 使用预定义的提示模板创建 SystemPromptTemplate
          SystemPromptTemplate systemPromptTemplate =
              new SystemPromptTemplate(REWRITE_SINGLE_QUERY_PROMPT);
          // 2. 使用 Map 填充模板中的占位符 {original_query}
          String systemMessageContent =
              systemPromptTemplate.createMessage(Map.of("original_query", originalQuery)).getText();
          // 3. 构建 Prompt。对于查询重写任务，通常将指令和原始查询都放入 SystemMessage。
          //    LLM 会根据 SystemMessage 的指令生成重写后的查询。
          Prompt prompt = new Prompt(List.of(new SystemMessage(systemMessageContent)));
          // 4. 调用 ChatClient 向 LLM 发送 Prompt 并获取响应
          String rewriteQuery = chatClient.prompt(prompt).call().content();
          System.out.println("Rewritten single query: " + rewriteQuery);
          // 5. 返回重写后的查询，如果 LLM 返回空或 null，则回退到原始查询
          return rewriteQuery != null ? rewriteQuery.trim() : originalQuery;
        },
        threadPoolTaskExecutor);
  }
}
